from flask import Flask, request, jsonify
import os
import json
from typing import Dict, List, Any, Optional
import argparse


class MCPServer:
    """
    A simple Model Context Protocol server that provides additional context
    based on user queries.
    """

    def __init__(self, data_dir: str = "./data"):
        """
        Initialize the MCP server.

        Args:
            data_dir: Directory containing context data
        """
        self.data_dir = data_dir
        self.knowledge_base = self._load_knowledge_base()

    def _load_knowledge_base(self) -> Dict[str, Any]:
        """
        Load the knowledge base from files in the data directory.

        In a real implementation, this might be a vector database,
        document store, or other more sophisticated system.

        Returns:
            A dictionary containing loaded knowledge
        """
        knowledge_base = {}

        # Create the data directory if it doesn't exist
        os.makedirs(self.data_dir, exist_ok=True)

        # Check for knowledge base file
        kb_path = os.path.join(self.data_dir, "knowledge_base.json")

        if os.path.exists(kb_path):
            try:
                with open(kb_path, 'r') as f:
                    knowledge_base = json.load(f)
            except json.JSONDecodeError:
                print(f"Error loading knowledge base from {kb_path}")
                knowledge_base = {}
        else:
            # Create a sample knowledge base if none exists
            sample_knowledge = {
                "model_context_protocol": {
                    "title": "Model Context Protocol (MCP)",
                    "description": "A protocol for providing additional context to language models.",
                    "details": [
                        "MCP enables models to access external data sources.",
                        "It provides a standardized way to augment model inputs with relevant context.",
                        "MCP servers can perform retrieval and processing before the model generates a response.",
                        "The protocol is designed to be model-agnostic and work with any LLM system."
                    ],
                    "use_cases": [
                        "Document retrieval",
                        "Knowledge base access",
                        "Database querying",
                        "API integration"
                    ]
                }
            }

            # Save the sample knowledge base
            with open(kb_path, 'w') as f:
                json.dump(sample_knowledge, f, indent=2)

            knowledge_base = sample_knowledge
            print(f"Created sample knowledge base at {kb_path}")

        return knowledge_base

    def process_query(self, query: str, history: Optional[List[Dict[str, str]]] = None) -> Dict[str, Any]:
        """
        Process a query and return relevant context.

        Args:
            query: The user query
            history: Optional chat history

        Returns:
            A dictionary with the relevant context and metadata
        """
        # Very simple keyword-based retrieval for demonstration
        # In a real implementation, this would use more sophisticated
        # retrieval techniques like embeddings and vector search

        context = ""
        context_used = False

        # Check for keywords in the query
        keywords = [
            "mcp", "model context protocol", "context protocol",
            "retrieve", "knowledge base", "external data"
        ]

        if any(keyword in query.lower() for keyword in keywords):
            # Found a match - retrieve the MCP information
            if "model_context_protocol" in self.knowledge_base:
                mcp_info = self.knowledge_base["model_context_protocol"]

                context = f"""
Title: {mcp_info['title']}
Description: {mcp_info['description']}

Details:
"""

                for detail in mcp_info['details']:
                    context += f"- {detail}\n"

                context += "\nUse Cases:\n"

                for use_case in mcp_info['use_cases']:
                    context += f"- {use_case}\n"

                context_used = True

        return {
            "context": context,
            "context_used": context_used
        }

    def add_to_knowledge_base(self, key: str, data: Dict[str, Any]) -> None:
        """
        Add new information to the knowledge base.

        Args:
            key: The key for the knowledge entry
            data: The data to store
        """
        self.knowledge_base[key] = data

        # Save the updated knowledge base
        kb_path = os.path.join(self.data_dir, "knowledge_base.json")
        with open(kb_path, 'w') as f:
            json.dump(self.knowledge_base, f, indent=2)


# Flask application
app = Flask(__name__)
mcp_server = None


@app.route('/health', methods=['GET'])
def health_check():
    """Health check endpoint."""
    return jsonify({"status": "healthy"})


@app.route('/query', methods=['POST'])
def process_query():
    """Process a query and return relevant context."""
    try:
        data = request.json

        if not data or 'query' not in data:
            return jsonify({"error": "Missing query parameter"}), 400

        query = data['query']
        history = data.get('history', [])

        result = mcp_server.process_query(query, history)

        return jsonify(result)

    except Exception as e:
        return jsonify({"error": str(e)}), 500


@app.route('/add', methods=['POST'])
def add_knowledge():
    """Add new knowledge to the knowledge base."""
    try:
        data = request.json

        if not data or 'key' not in data or 'data' not in data:
            return jsonify({"error": "Missing key or data parameters"}), 400

        key = data['key']
        knowledge_data = data['data']

        mcp_server.add_to_knowledge_base(key, knowledge_data)

        return jsonify({"status": "success", "message": f"Added knowledge for key: {key}"})

    except Exception as e:
        return jsonify({"error": str(e)}), 500


def main():
    """Run the MCP server."""
    parser = argparse.ArgumentParser(description="Model Context Protocol Server")
    parser.add_argument("--host", default="localhost", help="Host to run the server on")
    parser.add_argument("--port", default=3000, type=int, help="Port to run the server on")
    parser.add_argument("--data-dir", default="./data", help="Directory for knowledge base data")

    args = parser.parse_args()

    global mcp_server
    mcp_server = MCPServer(data_dir=args.data_dir)

    print(f"Starting MCP server on {args.host}:{args.port}")
    app.run(host=args.host, port=args.port)


if __name__ == "__main__":
    main()
